import math
import numpy as np
import control.matlab

import pat.autopilot as ap
import pat.vehicles.rotorcraft.dynamic_model_planar_simple as pdm
import pat.utils as pu

class Autopilot(ap.Autopilot):
    def __init__(self, _args=None, _dm=None, _time=None):
        ap.Autopilot.__init__(self)
        Xe, Ue = _dm.trim()
        for law in [DirectLaw, RateLaw, AttitudeLaw, AttitudeLawOpt, Attitude_z_Law_LQR,
                    VelocityLaw, PositionLaw, PositionLawOpt]:
            self.add_mode(law(_dm, Xe, Ue))
        self.set_cur_mode(6)

    def get_U_str(self):
        return "(f1, f2) {:-.3f} {:-.3f} N".format(self.U[pdm.DynamicModel.iv_f1], self.U[pdm.DynamicModel.iv_f2])
    

class Step_x(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, 0, "Step x") # x

class Step_z(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, 1, "Step z") # z

class Step_th(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, 2, "Step theta", period = 10., ampl=pu.rad_of_deg(10)) # theta



class DirectLaw(ap.LinearFeedback):
    def __init__(self, dm, Xe, Ue):
      K = np.array([[0., 0., 0., 0., 0., 0.],
                    [0., 0., 0., 0., 0., 0.]])
      ap.LinearFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, K, "Direct")
      


class RateLaw(ap.LinearFeedback):
    def __init__(self, dm, Xe, Ue):
      K = np.array([[0., 0., 0., 0., 0.,  0.2],
                    [0., 0., 0., 0., 0., -0.2]])
      ap.LinearFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, K, "Rate")



class AttitudeLaw(ap.LinearFeedback):
    def __init__(self, dm, Xe, Ue):
        Kp_th =  .05; Kd_th =  .025
        K = np.array([[0., 0.,  Kp_th, 0., 0.,  Kd_th],
                      [0., 0., -Kp_th, 0., 0., -Kd_th]])
        ap.LinearFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, K, "Attitude")



class AttitudeLawOpt(ap.LQRFeedback):
    def __init__(self, dm, Xe, Ue):
        A, B = dm.jacobian(Xe, Ue)
#        print A, B
        A1 = np.array([[A[dm.sv_th][dm.sv_th],  A[dm.sv_th][dm.sv_thd]],
                       [A[dm.sv_thd][dm.sv_th], A[dm.sv_thd][dm.sv_thd]]])
#        print A1
        B1 = np.array([[B[dm.sv_th][dm.iv_f1], B[dm.sv_th][dm.iv_f2]],
                       [B[dm.sv_thd][dm.iv_f1], B[dm.sv_thd][dm.iv_f2]]])
#        print B1
        #                 theta thetad
        self.Q = np.diag([0.01,  0.001])
        #                 f1   f2
        self.R = np.diag([0.1, 0.1])
        (K1, X1, E) = control.matlab.lqr(A1, B1, self.Q, self.R)
#        print K1
        K = np.array([[0., 0., -K1[dm.iv_f1][0], 0., 0., -K1[dm.iv_f1][1]],
                      [0., 0., -K1[dm.iv_f2][0], 0., 0., -K1[dm.iv_f2][1]]])
#        print K
        refs = [Step_th(Xe, Ue)]
        ap.LQRFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, K, "AttitudeOpt", refs)


class Attitude_z_Law_LQR(ap.LQRFeedback):
    def __init__(self, dm, Xe, Ue):
        A, B = dm.jacobian(Xe, Ue)
        K = np.array([[0., 0., 0., 0., 0., 0.],
                      [0., 0., 0., 0., 0., 0.]])
        refs = [Step_th(Xe, Ue)]
        ap.LQRFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, K, "Attitude z LQR", refs)


class VelocityLaw(ap.ControlLaw):
    def __init__(self, dm, Xe, Ue):
        ap.ControlLaw.__init__(self)
        self.name = "Velocity"


class PositionLaw(ap.LinearFeedback):
    def __init__(self, dm, Xe, Ue):
        self.refs = [Step_x(Xe, Ue), Step_z(Xe, Ue)]
        Kp_x  = -0.16 ;   Kd_x  = -0.155
        Kp_z  = -0.45;    Kd_z  = -0.68
        Kp_th =  0.60;    Kd_th =  0.105
        K = np.array([[ Kp_x, Kp_z,  Kp_th,  Kd_x, Kd_z,  Kd_th],
                      [-Kp_x, Kp_z, -Kp_th, -Kd_x, Kd_z, -Kd_th]])
        refs = [Step_x(Xe, Ue), Step_z(Xe, Ue)]
        ap.LinearFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, K, name="Position", refs=refs)
        

class PositionLawOpt(ap.LQRFeedback):
    def __init__(self, dm, Xe, Ue):
        A, B = dm.jacobian(Xe, Ue) 
        #                  x       z   theta  xd       zd     thetad
        self.Q = np.diag([0.005, 0.04, 0.05,  0.001,  0.002,  0.001])
        #                 f1   f2
        self.R = np.diag([0.05, 0.05])
        (K, X, E) = control.matlab.lqr(A, B, self.Q, self.R)
        refs = [Step_x(Xe, Ue), Step_z(Xe, Ue)]
        ap.LQRFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, -K, "PositionLQR", refs)

